import numpy as np
import matplotlib.pyplot as plt
import mpl_toolkits.axisartist as AA
import matplotlib.animation as animation

def gen_newtonForCos():
    # 使用牛顿方法解决优化问题
    # cosx = 0函数在(0, 3)区间的根
    # 迭代生成坐标
    x0 = 0.5  # 初始点位置<=pi
    for i in range(1000):
        f1order = -np.sin(x0)  # cos函数的一阶导数
        x1 = x0 - (np.cos(x0) / f1order)
        x0 = x1
        print("iteration ", i, ' x0: ', x0, ' x1: ', x1)

        if -0.0000000000001 < np.cos(x0) < 0.0000000000001:
            break

        yield x0


if __name__ == '__main__':

    x = np.linspace(0, 3)
    y = np.cos(x)

    fig = plt.figure()
    # 使用matplotlib工具包axisartist对axis进行处理
    ax = AA.Subplot(fig, 111)
    fig.add_axes(ax)
    ax.axis["right"].set_visible(False)
    ax.axis["top"].set_visible(False)
    ax.axis["bottom"].set_visible(False)
    ax.axis["left"].set_axisline_style("->", size=1.0)
    ax.axis["y=0"] = ax.new_floating_axis(nth_coord=0, value=0)
    ax.axis["y=0"].set_axisline_style("->", size=1.0)

    plt.title("newton method")
    plt.plot(x, y)
    plt.grid(False)

    point, = plt.plot(x[0], y[0], "ro")

    def update_points(x0):
        point.set_data(x0, np.cos(x0))
        local_x = np.arange(x0-2, x0+2)
        local_y = -np.sin(x0) * (local_x - x0) + np.cos(x0)
        plt.plot(local_x, local_y)

        if np.cos(x0) > 0:
            plt.vlines(x0, 0, np.cos(x0), color="green", linestyles="dashed")
        else:
            plt.vlines(x0, np.cos(x0), 0, color="yellow", linestyles="dashed")
        
        return point

    # 开始动画
    ani = animation.FuncAnimation(fig, update_points, gen_newtonForCos, interval=2000, repeat=False, blit=False)

    plt.show()
